import streamlit as st
from sqlalchemy import create_engine,inspect
from llama_index import SQLDatabase
import urllib
import openai
from url import *
import webbrowser

# openai
openai.api_key = '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'

st.set_page_config(page_title="Chatbot For Data", page_icon=":school:", menu_items={
    'Get Help': 'mailto:2301366947@qq.com',
    'Report a bug': "mailto:2301366947@qq.com",
    'About': '您可以提问有关excel和数据库的问题，本程序将给出您答案'
})

def generate_response_from_sql(query_engine, prompt):
    try:
        response = query_engine.query(prompt)

        sql = response.metadata['sql_query']
        sql_res = response.metadata['result']

        print(f"SQL语句为:\n{sql}")
        print(f"SQL查询结果为：{sql_res}")
        return str(response), sql
    except Exception as e:
        print(e)
        return "SQL无法找到答案", "", ""

def clear_chat_history():
    st.session_state.messages = [{"role": "assistant", "content": "我可以帮到您什么?"}]

st.write("<center><h2>Chat with data<h2></center>", unsafe_allow_html=True)

with st.sidebar:
    # 创建按钮
    if st.button("文档问答系统"):
        webbrowser.open_new_tab(qa_url)
    st.image('pku_logo.png')
    st.sidebar.markdown("## 数据库连接信息")
    username = st.text_input("username",placeholder="请输入用户名",key="username",value="root")
    password = st.text_input("password",placeholder="请输入密码",key="password",value="Zhaotianxiang74#",type="password")
    host = st.text_input("host",placeholder="请输入数据库地址",key="host",value="101.200.121.97:3306")
    database = st.text_input("database",placeholder="请输入数据库",key="database",value="reggie")


status = username and password and host and database

try:
    if status:
        encoded_password = urllib.parse.quote_plus(password)
        db_url = f'mysql+pymysql://{username}:{encoded_password}@{host}/{database}'
        engine = create_engine(db_url)

        # 创建数据库检查器对象
        inspector = inspect(engine)
        # 获取数据库中的所有表名
        tables = inspector.get_table_names()

        with st.sidebar.expander("数据库中的所有表如下"):
            st.write(tables)
        st.sidebar.button('清空对话历史', on_click=clear_chat_history)

        sql_database = SQLDatabase(engine, include_tables=tables)

        # text to sql
        from llama_index.indices.struct_store.sql_query import NLSQLTableQueryEngine

        query_engine = NLSQLTableQueryEngine(
            sql_database=sql_database,
            tables=tables,
            # synthesize_response=False,  # 只展示SQL查询结果
        )

        # Store LLM generated responses 存储LLM的响应信息
        if "messages" not in st.session_state.keys():
            st.session_state.messages = [{"role": "assistant", "content": "我可以帮到您什么?"}]

        # Display or clear chat messages
        for message in st.session_state.messages:
            with st.chat_message(message["role"]):
                st.write(message["content"])

        # User-provided prompt
        prompt = st.chat_input(placeholder="请输入问题...")
        if prompt:
            st.session_state.messages.append({"role": "user", "content": prompt})
            with st.chat_message("user"):
                st.write(prompt)

        if st.session_state.messages[-1]["role"] != "assistant":
            with st.chat_message("assistant"):
                with st.spinner("Querying..."):
                    response, sql = generate_response_from_sql(query_engine, prompt)
                    st.markdown("### SQL语句")
                    st.markdown(sql)
                    st.markdown("### 综合结果")
                    st.markdown(response)

            message = {"role": "assistant", "content": response}
            st.session_state.messages.append(message)
    else:
        st.header("请在左侧输入数据库连接信息")

except:
    st.header("数据库连接失败，请检查输入信息是否正确")